/*
 * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef TRT_DCNV2_PLUGIN_H
#define TRT_DCNV2_PLUGIN_H
#include "dcn_v2_im2col_cuda.h"
#include "kernel.h"
#include "plugin.h"

#include "serialize.hpp"
#include <cublas_v2.h>
#include <cuda.h>
#include <cudnn.h>
#include <string>
#include <vector>

using namespace nvinfer1::plugin;
namespace nvinfer1 {
namespace plugin {

class DeformableConvolutionalLayer : public IPluginV2Ext {
  public:
    DeformableConvolutionalLayer(int out_channels, int kernel, int deformable_group, int dilation, int padding,
                                 int stride, const Weights *weight, const Weights *bias);

    DeformableConvolutionalLayer(const void *buffer, size_t length);

    ~DeformableConvolutionalLayer() override = default;

    int getNbOutputs() const override;

    Dims getOutputDimensions(int index, const Dims *inputs, int nbInputs) override;

    int initialize() override;

    void terminate() override;

    size_t getWorkspaceSize(int maxBatchSize) const override;

    int enqueue(int batchSize, const void *const *inputs, void **outputs, void *workspace,
                cudaStream_t stream) override;

    size_t getSerializationSize() const override;

    void serialize(void *buffer) const override;

    bool supportsFormat(DataType type, PluginFormat format) const override;

    const char *getPluginType() const override;

    const char *getPluginVersion() const override;

    void destroy() override;

    IPluginV2Ext *clone() const override;

    void setPluginNamespace(const char *pluginNamespace) override;

    const char *getPluginNamespace() const override;

    DataType getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const override;

    bool isOutputBroadcastAcrossBatch(int outputIndex, const bool *inputIsBroadcasted, int nbInputs) const override;

    bool canBroadcastInputAcrossBatch(int inputIndex) const override;

    void attachToContext(cudnnContext *cudnnContext, cublasContext *cublasContext,
                         IGpuAllocator *gpuAllocator) override;

    void configurePlugin(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs,
                         const DataType *inputTypes, const DataType *outputTypes, const bool *inputIsBroadcast,
                         const bool *outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) override;

    void detachFromContext() override;

  private:
    Weights copyToDevice(const void *hostData, size_t count);
    void serializeFromDevice(char *&hostBuffer, Weights deviceWeights) const;
    Weights deserializeToDevice(const char *&hostBuffer, size_t count);

    std::string mPluginNamespace;

    int in_channels{};
    int height_out{};
    int width_out{};
    int height{};
    int width{};

    int out_channels{};
    int kernel_size{};
    int deformable_group{};
    int dilation{};
    int padding{};
    int stride{};

    Weights mWeight{};
    Weights mBias{};

    float *mOne;
    float *mColumn;

    cublasHandle_t mCublas;
};

class DCNv2PluginCreator : public BaseCreator {
  public:
    DCNv2PluginCreator();

    ~DCNv2PluginCreator() override = default;

    const char *getPluginName() const override;

    const char *getPluginVersion() const override;

    const PluginFieldCollection *getFieldNames() override;

    IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) override;

    IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) override;

  private:
    static PluginFieldCollection mFC;

    // Parameters for DeformableConvolutionalLayer
    static std::vector<PluginField> mPluginAttributes;
};
} // namespace plugin
} // namespace nvinfer1

#endif // TRT_DCNv2_PLUGIN_H
